#include "rknn_model_runner.hpp"

// OpenCV用于摄像头捕获
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc.hpp>

#include <sys/time.h>
#include <unistd.h>
#include <float.h>

// 获取当前时间（微秒）
static inline int64_t getCurrentTimeUs()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000000 + tv.tv_usec;
}

// 输出tensor属性信息
static void dump_tensor_attr(rknn_tensor_attr* attr)
{
    printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
        "zp=%d, scale=%f\n",
        attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
        attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
        get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

// 获取TopN预测结果
static int rknn_GetTopN(float* pfProb, float* pfMaxProb, uint32_t* pMaxClass, uint32_t outputCount, uint32_t topNum)
{
    uint32_t i, j;
    uint32_t top_count = outputCount > topNum ? topNum : outputCount;

    for (i = 0; i < topNum; ++i) {
        pfMaxProb[i] = -FLT_MAX;
        pMaxClass[i] = -1;
    }

    for (j = 0; j < top_count; j++) {
        for (i = 0; i < outputCount; i++) {
            if ((i == *(pMaxClass + 0)) || (i == *(pMaxClass + 1)) || (i == *(pMaxClass + 2)) || 
                (i == *(pMaxClass + 3)) || (i == *(pMaxClass + 4))) {
                continue;
            }

            if (pfProb[i] > *(pfMaxProb + j)) {
                *(pfMaxProb + j) = pfProb[i];
                *(pMaxClass + j) = i;
            }
        }
    }

    return 1;
}

// 从OpenCV Mat转换到RGB格式
int convert_opencv_mat_to_rgb(const cv::Mat& frame, uint8_t* out_buf, int* width, int* height)
{
    int ret = 0;

    if (frame.empty()) {
        printf("Input OpenCV Mat is empty!\n");
        return -1;
    }
    
    *width = frame.cols;
    *height = frame.rows;
    
    // 直接在输入Mat上进行原地转换
    if (frame.channels() == 1) {
        // 灰度图转RGB
        cv::Mat rgb_view(frame.rows, frame.cols, CV_8UC3, out_buf);
        cv::cvtColor(frame, rgb_view, cv::COLOR_GRAY2RGB);
    } else if (frame.channels() == 3) {
        // BGR转RGB (OpenCV默认是BGR顺序)
        cv::Mat rgb_view(frame.rows, frame.cols, CV_8UC3, out_buf);
        cv::cvtColor(frame, rgb_view, cv::COLOR_BGR2RGB);
    } else {
        printf("Unsupported OpenCV Mat format with %d channels!\n", frame.channels());
        return -1;
    }
    
    return ret;
}

// 构造函数
RknnModelRunner::RknnModelRunner(const char* model_path)
    : ctx(0), input_attrs(NULL), output_attrs(NULL),  output_mems(NULL), input_mems(NULL), 
    src_handle(0), is_initialized(false)
{
    // 初始化所有资源为0
    memset(&io_num, 0, sizeof(io_num));
    memset(&src_drm, 0, sizeof(src_drm));
    memset(&dst, 0, sizeof(dst));

    model_width = 0;
    model_height = 0;
    model_channel = 0;
    
    // 初始化模型
    if (init_model(model_path) != 0) {
        printf("Failed to initialize model %s\n", model_path);
        return;
    }
    
    // 分配模型输入输出内存
    if (allocate_io_memory() != 0) {
        printf("Failed to allocate IO memory\n");
        return;
    }
    
    is_initialized = true;
    printf("RknnModelRunner initialized successfully\n");
}

// 析构函数
RknnModelRunner::~RknnModelRunner()
{
    free_all_resources();
}

// 初始化模型
int RknnModelRunner::init_model(const char* model_path)
{
    int ret;
    
    // 1. 初始化RKNN上下文
    FILE* fp = fopen(model_path, "rb");
    if (fp == nullptr) {
        printf("fopen %s fail!\n", model_path);
        return -1;
    }
    fseek(fp, 0, SEEK_END);
    int            model_len = ftell(fp);
    unsigned char* model     = (unsigned char*)malloc(model_len);
    fseek(fp, 0, SEEK_SET);
    if (model_len != fread(model, 1, model_len, fp)) {
        printf("fread %s fail!\n", model_path);
        free(model);
        return -1;
    }
    if (fp) {
        fclose(fp);
    }

    ret = rknn_init(&ctx, model, model_len, 0, NULL);
    if (ret < 0) {
        printf("rknn_init fail! ret=%d\n", ret);
        return -1;
    }
    
    // rknn_core_mask core_mask = RKNN_NPU_CORE_0;
    // ret = rknn_set_core_mask(ctx, core_mask);

    // Get sdk and driver version
    rknn_sdk_version sdk_ver;
    ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
    if (ret != RKNN_SUCC) {
    printf("rknn_query fail! ret=%d\n", ret);
    return -1;
    }
    printf("rknn_api/rknnrt version: %s, driver version: %s\n", sdk_ver.api_version, sdk_ver.drv_version);

    // 2. 查询模型输入输出信息
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != RKNN_SUCC) {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("Model loaded: input_num: %d, output_num: %d\n", io_num.n_input, io_num.n_output);
    
    // 3. 分配内存存储模型输入输出属性
    input_attrs = (rknn_tensor_attr*)calloc(io_num.n_input, sizeof(rknn_tensor_attr));
    output_attrs = (rknn_tensor_attr*)calloc(io_num.n_output, sizeof(rknn_tensor_attr));
    if (!input_attrs || !output_attrs) {
        printf("Malloc tensor attrs failed\n");
        return -1;
    }
    
    // 4. 获取输入tensor属性
    printf("Input tensors:\n");
    for (uint32_t i = 0; i < io_num.n_input; i++) {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &input_attrs[i], sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            printf("rknn_query input_attr fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&input_attrs[i]);
    }
    
    // 5. 获取输出tensor属性
    printf("Output tensors:\n");
    for (uint32_t i = 0; i < io_num.n_output; i++) {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &output_attrs[i], sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            printf("rknn_query output_attr fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&output_attrs[i]);
    }
    
    // Get custom string
    rknn_custom_string custom_string;
    ret = rknn_query(ctx, RKNN_QUERY_CUSTOM_STRING, &custom_string, sizeof(custom_string));
    if (ret != RKNN_SUCC) {
    printf("rknn_query fail! ret=%d\n", ret);
    return -1;
    }
    printf("custom string: %s\n", custom_string.string);

    // 6. 获取模型输入尺寸
    if (input_attrs[0].fmt == RKNN_TENSOR_NHWC) {
        model_height = input_attrs[0].dims[1];
        model_width = input_attrs[0].dims[2];
        model_channel = input_attrs[0].dims[3];
    } else {
        model_height = input_attrs[0].dims[2];
        model_width = input_attrs[0].dims[3];
        model_channel = input_attrs[0].dims[1];
    }
    printf("Model required input size: %dx%dx%d\n", model_width, model_height, model_channel);
    
    return 0;
}

// 分配输入输出内存
int RknnModelRunner::allocate_io_memory()
{
    int ret;
    // 1. 分配RGA源图像DRM内存
    src = wrapbuffer_virtualaddr((void*)img_buf, CAMERA_WIDTH, CAMERA_HEIGHT, RK_FORMAT_RGB_888); // wstride, hstride,

    // 2. 分配NPU输入内存（RGA输出图像内存）
    input_attrs[0].type = RKNN_TENSOR_UINT8;
    input_attrs[0].fmt = RKNN_TENSOR_NHWC;
    input_mems = (rknn_tensor_mem**)calloc(1, sizeof(rknn_tensor_mem*));
    input_mems[0] = rknn_create_mem(ctx, input_attrs[0].size_with_stride);
    int wstride = model_width + (ALIGN - model_width % ALIGN) % ALIGN;
    int hstride = model_height;
    dst         = wrapbuffer_fd_t(input_mems[0]->fd, model_width, model_height, wstride, hstride,
                          RK_FORMAT_RGB_888); 

    // 3. 分配NPU输出内存
    output_mems = (rknn_tensor_mem**)calloc(io_num.n_output, sizeof(rknn_tensor_mem*));
    for (uint32_t i = 0; i < io_num.n_output; ++i) {
        // 使用float32类型输出，便于计算top5结果
        output_attrs[i].type = RKNN_TENSOR_FLOAT32;
        int output_size      = output_attrs[i].n_elems * sizeof(float);
        output_mems[i]       = rknn_create_mem(ctx, output_size);
    }
  
    // 4. 设置RKNN输入输出内存
    ret = rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]);
    if (ret < 0) {
        printf("rknn_set_io_mem for input failed! ret=%d\n", ret);
        return -1;
    }
    
    for (uint32_t i = 0; i < io_num.n_output; ++i) {
        ret = rknn_set_io_mem(ctx, output_mems[i], &output_attrs[i]);
        if (ret < 0) {
            printf("rknn_set_io_mem for output %d failed! ret=%d\n", i, ret);
            return -1;
        }
    }
    
    return 0;
}

// // 从图像文件进行推理
// int RknnModelRunner::inference_from_file(const char* image_path, int loop_count)
// {
//     if (!is_initialized) {
//         printf("Model not initialized!\n");
//         return -1;
//     }
    
//     int ret = 0;
//     int img_width = 0;
//     int img_height = 0;
//     int img_channel = 0;
    
//     // 加载图像
//     unsigned char* img_data = stbi_load(image_path, &img_width, &img_height, &img_channel, 4);
//     if (!img_data) {
//         printf("Load image %s failed!\n", image_path);
//         return -1;
//     }
//     printf("Load image: %s, w=%d, h=%d, c=%d\n", image_path, img_width, img_height, img_channel);
    
//     // 执行推理
//     ret = inference_from_image_data(img_data, img_width, img_height, loop_count);
    
//     // 释放图像内存
//     stbi_image_free(img_data);
    
//     return ret;
// }

// 从摄像头进行推理
int RknnModelRunner::inference_from_camera(int camera_index, int loop_count)
{
    if (!is_initialized) {
        printf("Model not initialized!\n");
        return -1;
    }
    
    int ret = 0;
    int img_width = 0;
    int img_height = 0;
    
    // 打开摄像头
    cv::VideoCapture cap(camera_index);
    if (!cap.isOpened()) {
        printf("Failed to open camera %d!\n", camera_index);
        return -1;
    }
    
    // 捕获一帧图像
    cv::Mat frame;
    printf("Capturing from camera %d...\n", camera_index);
    if (!cap.read(frame)) {
        printf("Failed to capture frame from camera!\n");
        return -1;
    }
    
    // 使用预分配的缓冲区转换为RGB格式
    ret = convert_opencv_mat_to_rgb(frame, img_buf, &img_width, &img_height);
    if (ret < 0) {
        printf("Failed to convert_opencv_mat_to_rgb!\n");
        return -1;
    }
    
    // 执行推理
    ret = inference_from_image_data(img_buf, img_width, img_height, loop_count);
    if (ret < 0) {
        printf("Failed to inference_from_image_data!\n");
        return -1;
    }
    return ret;
}

// 从RGBA图像数据进行推理
int RknnModelRunner::inference_from_image_data(uint8_t* img_data, int img_width, int img_height, int loop_count)
{
    if (!is_initialized || !img_data) {
        printf("Invalid parameters!\n");
        return -1;
    }
    
    int ret = 0;
    
    // 首先进行RGA缩放
    ret = do_rga_resize(img_data, img_width, img_height);
    if (ret != 0) {
        printf("Failed to do_rga_resize!\n");
        return ret;
    }
    
    // 然后执行推理
    ret = run_inference(loop_count);
    
    // 释放临时资源
    // free_temp_resources();
    
    return ret;
}

// 从物理连续内存执行RGA缩放
int RknnModelRunner::do_rga_resize(uint8_t* img_data, int img_width, int img_height)
{
    int ret = 0;
    
    // 1. 将原始图像数据拷贝到源DRM缓冲区
    // memcpy(src_drm.drm_buf, img_data, img_width * img_height * 3);
    // memcpy(dst_drm.drm_buf, 0, img_width * img_height * 3);

    // 2. 执行RGA缩放
    int64_t start_us, elapse_us;
    im_rect src_rect = {};
    im_rect dst_rect = {};
    
    // 检查RGA参数
    ret = imcheck(src, dst, src_rect, dst_rect);
    if (ret != IM_STATUS_NOERROR) {
        printf("RGA check error! %s\n", imStrError((IM_STATUS)ret));
        return -1;
    }
    
    // 执行RGA缩放
    printf("Performing RGA resize from %dx%d to %dx%d\n", img_width, img_height, model_width, model_height);
    start_us = getCurrentTimeUs();
    ret = imresize(src, dst);
    elapse_us = getCurrentTimeUs() - start_us;
    
    if (ret != IM_STATUS_SUCCESS) {
        printf("RGA resize failed! %s\n", imStrError((IM_STATUS)ret));
        return -1;
    }
    printf("RGA resize done in %.2fms\n", elapse_us / 1000.0f);
    
    return 0;
}

// 运行推理
int RknnModelRunner::run_inference(int loop_count)
{
    int ret = 0;
    int64_t start_us, elapse_us;
    
    // 执行推理
    printf("Begin inference (%d loops)...\n", loop_count);
    for (int i = 0; i < loop_count; ++i) {
        start_us = getCurrentTimeUs();
        ret = rknn_run(ctx, NULL);
        elapse_us = getCurrentTimeUs() - start_us;
        
        if (ret != RKNN_SUCC) {
            printf("rknn_run failed! ret=%d\n", ret);
            continue;
        }
        printf("Inference %d: Elapse time = %.2fms, FPS = %.2f\n", 
            i, elapse_us / 1000.0f, 1000.0f * 1000.0f / elapse_us);
    }
    
    // 处理推理结果
    uint32_t topNum = 5;
    for (uint32_t i = 0; i < io_num.n_output; ++i) {
        uint32_t MaxClass[topNum];
        float fMaxProb[topNum];
        float* buffer = (float*)output_mems[i]->virt_addr;
        uint32_t sz = output_attrs[i].n_elems;
        int top_count = sz > topNum ? topNum : sz;
        
        rknn_GetTopN(buffer, fMaxProb, MaxClass, sz, topNum);
        
        printf("---- Output %d - Top%d ----\n", i, top_count);
        for (int j = 0; j < top_count; ++j) {
            printf("%8.6f - Class %d\n", fMaxProb[j], MaxClass[j]);
        }
    }
    
    return 0;
}

// 释放所有资源
void RknnModelRunner::free_all_resources()
{
    // 1. 释放临时资源
    // free_temp_resources();
    
    
    // 2. 释放输入内存
    if (input_mems) {
        for (uint32_t i = 0; i < io_num.n_output; ++i) {
            if (input_mems[i]) {
                rknn_destroy_mem(ctx, input_mems[i]);
                input_mems[i] = NULL;
            }
        }
        free(input_mems);
        input_mems = NULL;
    }

    // 3. 释放输出内存
    if (output_mems) {
        for (uint32_t i = 0; i < io_num.n_output; ++i) {
            if (output_mems[i]) {
                rknn_destroy_mem(ctx, output_mems[i]);
                output_mems[i] = NULL;
            }
        }
        free(output_mems);
        output_mems = NULL;
    }
    
    // 5. 释放模型相关资源
    if (input_attrs) {
        free(input_attrs);
        input_attrs = NULL;
    }
    
    if (output_attrs) {
        free(output_attrs);
        output_attrs = NULL;
    }
    
    if (ctx) {
        rknn_destroy(ctx);
        ctx = 0;
    }
    
    if (src_handle > 0)
        releasebuffer_handle(src_handle);
    
    drm_buf_destroy(src_drm.drm_buffer_fd, src_drm.drm_buffer_handle, src_drm.drm_buf, src_drm.actual_size);

    is_initialized = false;
}

// 释放临时资源
// void RknnModelRunner::free_temp_resources()
// {
//     // 释放源图像资源
//     if (src_handle) {
//         releasebuffer_handle(src_handle);
//         src_handle = 0;
//     }
    
//     if (src_drm.drm_buf) {
//         drm_buf_destroy(src_drm.drm_buffer_fd, src_drm.drm_buffer_handle, src_drm.drm_buf, src_drm.actual_size);
//         src_drm.drm_buf = NULL;
//     }
// }
